# -*- coding: utf-8 -*-
# @日期    : 2021/11/28 13:49
# @作者  : 万方名
# @FileName: main.py

import numpy as np
import pandas as pd
import xgboost as xgb

data = pd.read_excel('../data/Concrete_Data.xls')

# 对label进行改名
data.rename(columns={'Concrete compressive strength(MPa, megapascals) ': 'label'}, inplace=True)

mask = np.random.rand(len(data)) < 0.8

train = data[mask]
test = data[~mask]

# 生成DMatrix
xgb_train = xgb.DMatrix(train.iloc[:, :7], label=train.label)
xgb_test = xgb.DMatrix(test.iloc[:, :7], label=test.label)

# train model
params = {
    'objective': 'reg:linear',  # 学习目标：线性回归
    'booster': 'gbtree',  # booster类型：树模型
    'min_child_weight': 1,  # 叶子结点最小样本权重和，若节点分裂导致叶子结点的样本权重和小于该值则节点不进行分裂
    'eta': 0.1,  # 学习率
    'max_depth': 5  # 决策树分裂的最大深度
}

num_round = 50
watchlist = [(xgb_train, 'train'), (xgb_test, 'test')]

model = xgb.train(params, xgb_train, num_round, watchlist)
model.save_model('../data/model.xgb')
